import numpy as np
import matplotlib.pyplot as plt

# 1.按要求完成下面的各项需求。
# 利用python编写如下程序，
# 现有一组猪肉销售数据集，包括训练集（d1.txt文件）和测试集（d2.txt文件）。数据集格式如下：
# 该数据集共有2列，第1列是猪肉重量，第2列是猪肉销售利润。
# 2.	请通过Python实现一元线性回归模型，并用此模型预测y，具体要求如下：
#
# 完成数据集的读取
train_data = np.loadtxt('./d1.txt', delimiter=',')
print(train_data[:5])
m = len(train_data)
x = train_data[:, 0:1]
XX = np.c_[np.ones([m, 1]), x]
y = train_data[:, -1]

test_data = np.loadtxt('./d2.txt', delimiter=',')
print(test_data[:5])
m_test = len(test_data)
x_test = test_data[:, 0:1]
XX_test = np.c_[np.ones([m_test, 1]), x_test]
y_test = test_data[:, -1]

# 实现线性回归的代价函数
def model(X, theta):
    return X.dot(theta)


def cost_func(h, y):
    m = len(h)
    e = h - y
    sq = e.dot(e)
    j = 1 / (2 * m) * sq
    return j


# 实现梯度下降函数
def grad(X, y, theta=None, alpha=0.001, num_iter=15000):
    m, n = X.shape
    if theta is None:
        theta = np.zeros(n)
    j_his = np.zeros(num_iter)

    for i in range(num_iter):
        h = model(X, theta)
        j = cost_func(h, y)
        j_his[i] = j
        e = h - y
        dt = 1.0 / m * X.T.dot(e)
        theta -= alpha * dt
    return theta, j_his, h


theta, j_his_train, h_train = grad(XX, y)
print(f'Theta = {theta}')

#     要求输出迭代过程中的代价函数值
plt.figure(figsize=[12, 5])
spr = 1
spc = 2
spn = 0
spn += 1
plt.subplot(spr, spc, spn)
plt.plot(j_his_train, label='cost function value')
plt.grid()
plt.legend()

# 通过梯度下降计算回归模型，用所得模型对测试集的数据进行预测
h_test = model(XX_test, theta)

# 以横轴为真实值，纵轴为预测值，画出散点图进行对比
spn += 1
plt.subplot(spr, spc, spn)
plt.scatter(y_test, y_test, s=1, label='target')
plt.scatter(y_test, h_test, s=1,label='hypothesis')
plt.grid()
plt.legend()

plt.show()
